package com.kafka.study.consumer;

import jakarta.annotation.PreDestroy;
import jakarta.annotation.Resource;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.Consumer;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.OffsetAndTimestamp;
import org.apache.kafka.common.PartitionInfo;
import org.apache.kafka.common.TopicPartition;
import org.springframework.beans.factory.annotation.Qualifier;
import org.springframework.beans.factory.annotation.Value;
import org.springframework.kafka.core.ConsumerFactory;
import org.springframework.stereotype.Service;

import java.time.Duration;
import java.time.Instant;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.concurrent.RejectedExecutionException;
import java.util.concurrent.ThreadPoolExecutor;
import java.util.stream.Collectors;

@Service
@Slf4j
public class ScheduleConsumer {
    // consumer不是线程安全的，在多线程的情况下需要用锁机制
    private final Consumer<String, String> consumer;

    public ScheduleConsumer(@Qualifier("timeConsumerFactory") ConsumerFactory<String, String> factory) {
        consumer = factory.createConsumer();
    }

    // TODO 自定义线程池
    @Resource
    private ThreadPoolExecutor executor;

    @Value("${spring.kafka.template.time-topic}")
    private String timeTopic;

    @Value("${spring.kafka.consumer.time.time-window}")
    private int timeWindow;

    /**
     * 定时调度
     */
    public synchronized void consume() {
        try {
            List<PartitionInfo> partitionInfos = consumer.partitionsFor(timeTopic);
            List<TopicPartition> topicPartitions = partitionInfos.stream().map(partitionInfo -> new TopicPartition(timeTopic, partitionInfo.partition())).toList();
            consumer.assign(topicPartitions);

            Instant now = Instant.now();
            long start = now.minusSeconds(timeWindow).toEpochMilli();
            long end = now.plusSeconds(timeWindow).toEpochMilli();
            Map<TopicPartition, Long> timestampToSearch = consumer.assignment().stream().collect(Collectors.toMap(tp -> tp, tp -> start));
            Map<TopicPartition, OffsetAndTimestamp> offsets = consumer.offsetsForTimes(timestampToSearch);
            for (TopicPartition tp : consumer.assignment()) {
                OffsetAndTimestamp offsetAndTimestamp = offsets.get(tp);
                if (offsetAndTimestamp != null) {
                    consumer.seek(tp, offsetAndTimestamp.offset());
                } else {
                    consumer.seekToEnd(Collections.singletonList(tp));
                }
            }

            ConsumerRecords<String, String> records = consumer.poll(Duration.ofSeconds(3));
            if (records.isEmpty()) {
                return;
            }
            records.forEach(record -> {
                try {
                    executor.submit(() -> {
                        handleBusiness(record, start, end);
                    });
                } catch (RejectedExecutionException e) {
                    handleException(record, e);
                }
            });
        } catch (Exception e) {
            log.error("error");
            throw new RuntimeException(e);
        }
    }

    private void handleException(ConsumerRecord<String, String> record, RejectedExecutionException e) {
        // 重试或者告警
    }

    private void handleBusiness(ConsumerRecord<String, String> record, long start, long end) {
        if (record.timestamp() < start) {
            log.warn("消息的时间已经过期了");
            return;
        }
        if (record.timestamp() > end) {
            log.warn("消息的时间还没到");
            return;
        }
        // 在时间范围内的数据进行业务处理
        // 可以用redis分布式锁保证幂等性
        log.info("处理业务");
    }

    /**
     * 防止资源泄露
     */
    @PreDestroy
    public void cleanup() {
        if (consumer != null) {
            consumer.close();
        }
    }
}
